package top.jolyoulu.core.acc

import org.apache.spark.rdd.RDD
import org.apache.spark.util.{AccumulatorV2, LongAccumulator}
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable

/**
 * @Author: JolyouLu
 * @Date: 2024/2/6 16:32
 * @Description
 */
object Spark01_RDD_Acc_WordCount {
  def main(args: Array[String]): Unit = {
    val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("Acc")
    val sc: SparkContext = new SparkContext(sparkConf)
    //创建累加器对象
    val wcAcc: MyAccumulator = new MyAccumulator
    //注册累加器
    sc.register(wcAcc, "wordCountAcc")

    val rdd: RDD[String] = sc.makeRDD(List("hello", "spark", "hello"))

    rdd.foreach(word => {
      wcAcc.add(word)
    })

    println(wcAcc.value)

    //关闭环境
    sc.stop()
  }

  /**
   * 自定义累加器
   * 继承AccumulatorV2，定义泛型
   * IN：累加器输入的数据类型
   * OUT：累加器返回的数据类型
   */
  class MyAccumulator extends AccumulatorV2[String, mutable.Map[String, Long]] {

    private var wcMap = mutable.Map[String, Long]()

    //判断是否初始状态
    override def isZero: Boolean = {
      wcMap.isEmpty
    }

    override def copy(): AccumulatorV2[String, mutable.Map[String, Long]] = {
      new MyAccumulator()
    }

    override def reset(): Unit = {
      wcMap.clear()
    }

    //获取累加器需要计算的值
    override def add(word: String): Unit = {
      val newCnt: Long = wcMap.getOrElse(word, 0L) + 1
      wcMap.update(word, newCnt)
    }

    //各Executor的累加器累加后，最后会全部merge到Driver
    override def merge(other: AccumulatorV2[String, mutable.Map[String, Long]]): Unit = {
      val map1 = this.wcMap
      val map2 = other.value
      map2.foreach {
        case (word, count) => {
          val newCount = map1.getOrElse(word, 0L) + count
          map1.update(word, newCount)
        }
      }
    }

    override def value: mutable.Map[String, Long] = {
      wcMap
    }
  }
}
